import torch
import torch.nn as nn

class C3D(nn.Module):
    def __init__(self, num_classes,pretrained=False):
        super(C3D, self).__init__()
        self.conv1 = nn.Conv3d(3, 64, (3,3,3),  padding=(1,1,1))#卷积层
        self.pool1 = nn.MaxPool3d(kernel_size=(1,2,2), stride=(1,2,2))#池化层
        #第2块
        self.conv2 = nn.Conv3d(64, 128, (3, 3, 3), padding=(1, 1, 1))  # 卷积层
        self.pool2= nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2))  # 池化层
        #第3块
        self.conv3a= nn.Conv3d(128, 256, (3, 3, 3), padding=(1, 1, 1))  # 卷积层
        self.conv3b= nn.Conv3d(256, 256, (3, 3, 3), padding=(1, 1, 1))  # 卷积
        self.pool3= nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2))  # 池化层
        #第4块
        self.conv4a= nn.Conv3d(256, 512, (3, 3, 3), padding=(1, 1, 1))
        self.conv4b= nn.Conv3d(512, 512, (3, 3, 3), padding=(1, 1, 1))
        self.pool4= nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2))
        #第5块
        self.conv5a= nn.Conv3d(512, 512, (3, 3, 3), padding=(1, 1, 1))
        self.conv5b= nn.Conv3d(512, 512, (3, 3, 3), padding=(1, 1, 1))
        self.pool5= nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2), padding=(0,1,1))
        #全连接层
        self.fc6=nn.Linear(8192, 4096)
        self.fc7=nn.Linear(4096, 4096)
        self.fc8=nn.Linear(4096, num_classes)

        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)#定义Dropout层
        self.__init_weights()#初始化
        self.__init_weights()

        if pretrained:#可以选择是否加载预训练权重，默认不加载
            self.__load_pretrained_model()
    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.pool1(x)

        x = self.relu(self.conv2(x))
        x = self.pool2(x)

        x = self.relu(self.conv3a(x))
        x = self.relu(self.conv3b(x))
        x = self.pool3(x)

        x = self.relu(self.conv4a(x))
        x = self.relu(self.conv4b(x))
        x = self.pool4(x)

        x = self.relu(self.conv5a(x))
        x = self.relu(self.conv5b(x))
        x = self.pool5(x)
        #全连接层
        x = x.view(-1,8192)  # x = x.view(x.size(0), -1)，二者一样，x.size(0)是批次数量
        x=self.relu(self.fc6(x))
        x=self.dropout(x)
        x=self.relu(self.fc7(x))
        x=self.dropout(x)

        x=self.fc8(x)
        return x


    #权重参数初始化
    def __init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv3d):#判断网络层是不是conv3d
                nn.init.kaiming_normal_(m.weight)

    #加载预训练权重
    def __load_pretrained_model(self):
        corresp_name={
            #Conv1-conv7的权重和偏重名字，已经用print(name)打印出。最后一层分布不一样，不能直接平替
        }
        p_dict = torch.load("ucf101-caffe.pth")#已经训练好的
        s_dict = self.state_dict()#加载自己模型参数
        for name in p_dict():
            #print(name)#打印别人模型的名字
            if name not in corresp_name:#不在该字典里，跳过这个
                continue
            s_dict[corresp_name[name]] = p_dict[name]#等价于s_dict['conv1.weight']
            self.load_state_dict(s_dict)#重新加载一下复制的值



if __name__ == '__main__':
    from torchsummary import summary
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    net=C3D(num_classes=101).to(device)
    print(summary(net,(3,16,112,112)))

